import json
import os

labels = ["pig", "person"]
def labelme_to_yolo_segment(labelme_json_path, output_dir):
    with open(labelme_json_path, "r") as f:
        data = json.load(f)

    image_height = data["imageHeight"]
    image_width = data["imageWidth"]
    yolo_label = ""
    for shape in data["shapes"]:
        label = shape["label"]
        if label not in labels:
            continue
        points = shape["points"]


        # 将多边形点转换为 YOLO 格式
        yolo_format_points = []
        for point in points:
            x = point[0] / image_width
            y = point[1] / image_height
            yolo_format_points.append(f"{x:.6f} {y:.6f}")

        cidx = labels.index(label)
        yolo_label += f"{cidx} " + " ".join(yolo_format_points) + "\n"

        # 保存为 YOLO 格式的标签文件
    output_file_path = os.path.join(
        output_dir, os.path.splitext(os.path.basename(labelme_json_path))[0] + ".txt"
    )
    with open(output_file_path, "w") as out_f:
        out_f.write(yolo_label + "\n")


def convert_all_json_in_directory(input_dir, output_dir):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for filename in os.listdir(input_dir):
        if filename.endswith(".json"):
            labelme_json_path = os.path.join(input_dir, filename)
            labelme_to_yolo_segment(labelme_json_path, output_dir)


# 示例用法
input_dir = "C:/Users/kang_/Desktop/d"
output_dir = "C:/Users/kang_/Desktop/d"
convert_all_json_in_directory(input_dir, output_dir)
